package com.rpcservice.rpc.loadBalance.loadbalancer;

import com.rpcservice.rpc.loadBalance.AbstractLoadBalance;
import com.rpcservice.rpc.remoting.dto.RpcRequest;
import lombok.extern.java.Log;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentHashMap;

/**
 * refer to dubbo consistent hash load balance:
 * dubbo的一致性哈希负载均衡算法实现参考：
 * https://github.com/apache/dubbo/blob/2d9583adf26a2d8bd6fb646243a9fe80a77e65d5/dubbo-cluster
 * /src/main/java/org/apache/dubbo/rpc/cluster/loadbalance/ConsistentHashLoadBalance.java
 * *
 */
public class ConsistentHashLoadBalance extends AbstractLoadBalance {
    private static final Logger logger = LoggerFactory.getLogger(ConsistentHashLoadBalance.class);

    private final ConcurrentHashMap<String, ConsistentHashSelector> selectors = new ConcurrentHashMap<>();

    /**
     * 根据服务地址列表和RPC请求选择一个服务地址
     *
     * @param serviceUrlList
     * @param rpcRequest
     * @return
     */
    @Override
    protected String doSelect(List<String> serviceUrlList, RpcRequest rpcRequest) {
        //获取调用者列表的身份哈希码
        int identityHashCode = System.identityHashCode(serviceUrlList);
        //获取RPC服务名称
        String rpcServiceName = rpcRequest.getRpcServiceName();
        //从选择器映射中获取对应的选择器
        ConsistentHashSelector selector = selectors.get(rpcServiceName);
        //如果选择器不存在或身份哈希码不匹配，则创建一个新的选择器
        if (selector == null || selector.identityHashCode != identityHashCode) {
            selectors.put(rpcServiceName, new ConsistentHashSelector(serviceUrlList, 160, identityHashCode));
            selector = selectors.get(rpcServiceName);
        }
        //使用选择器根据RPC请求选择一个服务地址, 参数是rpcServiceName加上请求参数的数组字符串，Arrays.stream可以将参数数组转换为流
        //比如rpcServiceName为"com.example.MyService"，参数为["param1", "param2"]，则传入的字符串为"com.example.MyService[param1, param2]"
        return selector.select(rpcServiceName + Arrays.stream(rpcRequest.getParameters()));
    }


    /**
     * 一致性哈希选择器的内部类
     */
    static class ConsistentHashSelector {
        /**
         * 虚拟节点的哈希环
         */
        private final TreeMap<Long, String> virtualInvokers;
        /**
         * identity hash code of the invokers list
         * 用于标识调用者列表的身份哈希码
         */
        private final int identityHashCode;

        ConsistentHashSelector(List<String> invokers, int replicaNumber, int identityHashCode) {
            this.virtualInvokers = new TreeMap<>();
            this.identityHashCode = identityHashCode;
            //这是遍历调用者列表，并为每个调用者创建多个虚拟节点
            for (String invoker : invokers) {
                for (int i = 0; i < replicaNumber / 4; i++) {
                    byte[] digest = md5(invoker + i);
                    for (int h = 0; h < 4; h++) {
                        long m = hash(digest, h);
                        virtualInvokers.put(m, invoker);
                    }
                }
            }
        }

        /**
         * md5 hash算法,原理是
         * 1.将输入数据进行分块处理，每个块的大小为512位（64字节）。
         * 2.对每个数据块进行一系列的位运算和非
         * 线性函数处理，生成一个128位（16字节）的哈希值。
         * 3.将所有数据块的哈希值进行组合，得到最终的哈希值。
         *
         * @param key
         * @return
         */
        static byte[] md5(String key) {
            MessageDigest md = null;
            try {
                md = MessageDigest.getInstance("MD5");
                byte[] bytes = key.getBytes(StandardCharsets.UTF_8);
                md.update(bytes);
            } catch (NoSuchAlgorithmException e) {
                //请求特定的加密算法不可用时抛出异常
                throw new IllegalStateException(e.getMessage(), e);
            }
            return md.digest();
        }

        /**
         * 计算哈希值
         * 计算公式是将字节数组中的四个字节组合成一个32位的整数，
         * 并将其转换为无符号长整型。
         *
         * @param digest
         * @param idx
         * @return
         */
        static long hash(byte[] digest, int idx) {
            return (((long) (digest[3 + idx * 4] & 0xFF) << 24)
                    | ((long) (digest[2 + idx * 4] & 0xFF) << 16)
                    | ((long) (digest[1 + idx * 4] & 0xFF) << 8)
                    | ((long) (digest[0 + idx * 4] & 0xFF))) & 0xFFFFFFFFL;
        }

        /**
         * 根据rpcServiceKey选择对应的虚拟节点
         *
         * @param rpcServiceKey
         * @return
         */
        public String select(String rpcServiceKey) {
            //先MD5计算哈希值
            byte[] digest = md5(rpcServiceKey);
            //然后在hash()方法中计算出对应的哈希值,得到无符号长整型
            long hashCode = hash(digest, 0);
            return selectForKey(hashCode);
        }

        /**
         * 根据哈希值选择对应的虚拟节点
         *
         * @param hashCode
         * @return
         */
        public String selectForKey(long hashCode) {
            //ceilingEntry方法返回大于或等于给定键的最小键值对
            Map.Entry<Long, String> entry = virtualInvokers.ceilingEntry(hashCode);
            if (entry == null) {
                entry = virtualInvokers.firstEntry();
            }
            return entry.getValue();
        }
    }
}
